import os
from sklearn.model_selection import train_test_split
import cv2
import numpy as np
import pandas as pd
from utils.rle_parse import make_mask
import json


def provider(
        data_folder,
        df_path,
        save_folder='./temp',
):
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)
        os.makedirs(os.path.join(save_folder, 'train'))
        os.makedirs(os.path.join(save_folder, 'test'))
        os.makedirs(os.path.join(save_folder, 'test', 'ok'))
        os.makedirs(os.path.join(save_folder, 'test', 'ko'))
        os.makedirs(os.path.join(save_folder, 'ground_truth'))
        os.makedirs(os.path.join(save_folder, 'ground_truth', 'ko'))

    df = pd.read_csv(df_path)
    df['ImageId'], df['ClassId'] = zip(*df['ImageId_ClassId'].str.split('_'))
    df['ClassId'] = df['ClassId'].astype(int)
    df = df.pivot(index='ImageId', columns='ClassId', values='EncodedPixels')
    df['defects'] = df.count(axis=1)
    error_num = 0
    normal_num = 0
    normal_list = []
    for i in range(len(df)):
        image_id, mask = make_mask(i, df)
        mask = np.sum(mask, axis=2)
        mask[mask > 0] = 255
        image_path = os.path.join(data_folder, "train_images", image_id)
        image = cv2.imread(image_path)
        image = cv2.resize(image, (1536, 256))
        mask = cv2.resize(mask, (1536, 256))
        image_id = image_id.split('.')[0]
        for x in range(6):
            cropped_image = image[:, x * 256:x * 256 + 256]
            cropped_mask = mask[:, x * 256:x * 256 + 256]
            if cropped_mask.sum() != 0:
                error_num += 1
                path_img = os.path.join(save_folder, 'test', 'ko', f'{image_id}_index_{x}.png')
                path_mask = os.path.join(save_folder, 'ground_truth', 'ko', f'{image_id}_index_{x}.png')
                cv2.imwrite(path_img, cropped_image)
                cv2.imwrite(path_mask, cropped_mask)
            if cropped_mask.sum() == 0:
                normal_num += 1
                normal_list.append({'id': f'{image_id}_index_{x}.png', 'image': cropped_image})
    part_one, part_two = train_test_split(normal_list, test_size=0.16, random_state=42)
    for item in part_one:
        image_path = os.path.join(save_folder, "train", item['id'])
        cv2.imwrite(image_path, item['image'])
    for item in part_two:
        image_path = os.path.join(save_folder, "test", "ok", item['id'])
        cv2.imwrite(image_path, item['image'])
    print("有缺陷", error_num)
    print("无缺陷", normal_num)


class SteelSolver(object):

    def __init__(self, root='steel'):
        self.root = root
        self.meta_path = f'{root}/meta.json'
        self.phases = ['train', 'test']

    def run(self):
        info = {phase: {} for phase in self.phases}
        # cls_dir = f'{self.root}/steel'
        for phase in self.phases:
            cls_info = []
            ok_ko = os.listdir(f'{self.root}/{phase}')
            for c in ok_ko:
                is_abnormal = True if c not in ['ok'] else False

                img_dir = f'{self.root}/{phase}/{c}'
                mask_dir = f'{self.root}/ground_truth/{c}'
                img_names = os.listdir(img_dir)
                mask_names = os.listdir(mask_dir) if is_abnormal else None
                img_names.sort()
                mask_names.sort() if mask_names is not None else None
                for idx, img_name in enumerate(img_names):
                    info_img = dict(
                        img_path=f'{img_dir.replace(self.root, "")}/{img_name}',
                        mask_path=f'{mask_dir.replace(self.root, "")}/{mask_names[idx]}' if is_abnormal else '',
                        cls_name='steel',
                        specie_name='steel',
                        anomaly=1 if is_abnormal else 0,
                    )
                    cls_info.append(info_img)
            info[phase]['steel'] = cls_info
        with open(self.meta_path, 'w') as f:
            f.write(json.dumps(info, indent=4) + "\n")


if __name__ == '__main__':
    # provider(data_folder='./Steel_data', df_path='./Steel_data/train.csv', save_folder='./temp')
    solver = SteelSolver()
    solver.run()
